# -*- coding: utf-8 -*-
# function to import mathematica simulation result

from scipy import *
from scipy import interpolate
import pickle

def stable_check(A, theta):
    # Check for bistable state for SW particle
    # A hack to revert the theta value into the 0 to pi/2 range.
    theta = abs(pi*floor(theta/(pi/2))-theta)
    t = tan(theta)**(1./3.)
    t2=t*t
    t4=t2*t2
    front = sqrt(1+t2+t4)/(1+t2)
    return A<front

def const_interpolate(fname):
    '''Load the interpolation data for further use'''
    print('Load interpolation data file %s' %fname)
    int_data = loadtxt(fname)
    alist = int_data[:,0]
    tlist = int_data[:,1]
    p1list = int_data[:,2]
    p2list = int_data[:,3]
    e1list = int_data[:,4]
    e2list = int_data[:,5]
    e3list = int_data[:,6]
    #test_alist = arange(0,0.99,0.04)
    #test_tlist = arange(0,pi/2+0.00001,pi/40)
    #test_p1list = p1list.copy()
    #test_p1list.resize((len(test_alist),len(test_tlist)))
    #test_p2list = p2list.copy()
    #test_p2list.resize((len(test_alist),len(test_tlist)))
    
    coord = stable_check(alist,tlist)
    select_alist=alist[coord]
    select_tlist=tlist[coord]
    select_p2list=p2list[coord]
    select_e2list=e2list[coord]
    
    p1=interpolate.interp2d(alist,tlist,p1list,kind='cubic')
    p2=interpolate.interp2d(select_alist,select_tlist,select_p2list,kind='cubic')
    e1=interpolate.interp2d(alist,tlist,e1list,kind='cubic')
    e2=interpolate.interp2d(select_alist,select_tlist,select_e2list,kind='cubic')
    e3=interpolate.interp2d(alist,tlist,e3list,kind='cubic')
    
    outnames=[]
    
    for fun in [(p1,'p1'),(p2,'p2'),(e1,'e1'),(e2,'e2'),(e3,'e3')]:
        outname =  fun[1]+'.pickle'
        with open(outname,'w') as fname:
            pickle.dump(fun[0].tck, fname)
            outnames.append(outname)
            print('Pickle the interpolation function %s.' %outname)
            
    return outnames

if __name__=='__main__':
    int_filename='SW.dat'
    # we start to load the data directly
    outnames = const_interpolate(int_filename)